Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@deepmd/dpmodel/utils/dist_check.py`:
- Around line 45-56: The current implementation materializes diff =
real_coord[np.newaxis, :, :] - real_coord[:, np.newaxis, :] which creates an
(n_real, n_real, 3) array and O(n^2) memory; replace this with a blocked (or
streaming) pairwise-min computation to avoid full NxN allocation: iterate over
chunks/blocks of indices (e.g., outer blocks i and inner blocks j), compute the
block-wise displacement between real_coord[i:i_b] and real_coord[j:j_b], apply
the PBC transform using cell and inv_cell (reuse the same frac_diff -=
round(frac_diff) and back-transform logic on the block), compute block-wise
dist_sq, ignore self-pair entries when i==j (mask diagonal within the block),
and accumulate the global minimum distance (or a per-atom min array) instead of
assembling dist_sq for all pairs; update code locations around symbols diff,
frac_diff, cell, inv_cell, and dist_sq to perform block-wise operations and
maintain identical results with bounded memory.

In `@deepmd/pt/entrypoints/freeze_pt2.py`:
- Around line 377-380: The type check in freeze_sezm_to_pt2 currently only
allows "sezm" and rejects valid "dpa4" checkpoints; update the validation to
accept both "sezm" and "dpa4" (case-insensitive) by checking params.get("type")
against the set {"sezm", "dpa4"} so that is_sezm_checkpoint()-routed dpa4
checkpoints are not erroneously rejected.

In `@deepmd/pt/model/descriptor/sezm_nn/indexing.py`:
- Line 70: Rename the ambiguous single-letter identifier `l` to a clearer name
such as `ell` throughout indexing.py to satisfy Ruff E741; specifically update
list comprehensions like `[2 * l + 1 for l in range(lmax + 1)]` to `[2 * ell + 1
for ell in range(lmax + 1)]`, and rename the loop/parameter occurrences
referenced in the file (including the other reported occurrences around the same
blocks) so function signatures, list comprehensions, for-loops and any uses
inside functions/methods (e.g., variables used alongside `lmax`) are
consistently changed to `ell` (and all references updated) to avoid breaking
references. Ensure tests/uses of symbols from functions/methods in this module
reflect the new name and run `ruff check .` to confirm E741 is resolved.

In `@deepmd/pt/model/descriptor/sezm_nn/radial.py`:
- Around line 230-241: The constructor currently allows non-positive rcut which
later causes invalid divisions; in the class __init__ (the shown __init__) and
the other constructor/initializer that assigns self.rcut (the occurrence around
lines 487-488), add a strict check that rcut > 0 and raise ValueError("`rcut`
must be positive") if not, then convert rcut to float and assign to self.rcut
(and only derive widths or compute 1/rcut after that check) so all downstream
computations are safe.

---

Nitpick comments:
In `@deepmd/pt/model/descriptor/sezm_nn/so2.py`:
- Around line 1340-1342: The zip over the three module lists (so2_linears,
so2_inter_norms, non_linearities) should use strict=True to surface any length
mismatches at runtime; update the zip(...) calls in the for loops that iterate
over (so2_linear, inter_norm, non_linear) (and the similar loop around lines
1364–1366) to call zip(self.so2_linears, self.so2_inter_norms,
self.non_linearities, strict=True) so a ValueError is raised if the lists differ
in length.

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@deepmd/pt/model/descriptor/sezm_nn/so2.py`:
- Around line 1341-1343: The two for-loops that iterate with
zip(self.so2_linears, self.so2_inter_norms, self.non_linearities) should use
zip(..., strict=True) to enforce equal-length sequences at runtime; update both
occurrences (the loop that assigns layer_idx, so2_linear, inter_norm, non_linear
and the later similar loop) so they pass strict=True to zip to catch future
mismatches defensively.

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@deepmd/pt/entrypoints/freeze_pt2.py`:
- Around line 268-285: ntypes should be derived from the model's actual reported
type count rather than len(model.get_type_map()), because get_type_map() may be
empty; replace the ntypes assignment so it queries the model for its canonical
type count (e.g., model.get_type_count() or equivalent API) and fall back to
max(1, len(model.get_type_map())) if that method is unavailable, then use that
ntypes when populating atype_np in the loop that assigns i % ntypes.

In `@deepmd/pt/model/descriptor/sezm_nn/so2.py`:
- Around line 1341-1366: The two for-loops that iterate over
zip(self.so2_linears, self.so2_inter_norms, self.non_linearities) raise Ruff
B905; update both zips to use strict=True (i.e., zip(..., strict=True)) so
mismatched lengths will error and the linter warning is resolved — change the
zip in the first loop (around so2_layer_attn_res usage and so2_depth_sources
appends) and the zip in the else branch to include strict=True while leaving the
rest of the loop bodies unchanged.
- Around line 324-329: The cached tensor self._cached_weight can become stale or
on the wrong device after module.to(...) or load_state_dict(); update the class
to invalidate this cache on device/dtype moves and state loads by overriding
nn.Module._apply(self, fn) to set self._cached_weight = None (so moves clear the
cache) and/or override _load_from_state_dict(...) to clear self._cached_weight
when parameters are loaded; alternatively (or in addition) add a runtime check
where the cache is used (the weight-access/forward path that reads
self._cached_weight) to verify tensor.device and tensor.dtype match the module
parameters before reusing, and if not, set self._cached_weight = None and
rebuild.

In `@deepmd/pt/model/descriptor/sezm.py`:
- Around line 1937-1941: The current loop that strips every unexpected key under
the descriptor prefix (using prefix, state_dict(), expected_keys and iterating
over state_dict.keys()) must be replaced with a targeted whitelist of only known
transient keys that are intentionally rebuilt at construction; stop
blanket-popping unknown keys so load_state_dict(strict=True) can report
unexpected_keys. Update the logic to define an explicit set/list of transient
keys (or a clear pattern) that you will remove (e.g., transient_keys = {...} or
check exact suffixes), only pop entries that match those transient keys under
prefix, and then call load_state_dict(strict=True) so any other mismatches
surface as errors. Ensure references to prefix and expected_keys remain correct
when computing which keys are transient vs genuine parameters.
- Around line 1137-1152: forward_with_edges bypasses forward's normalization of
charge_spin, causing _apply_charge_spin_embedding to receive uncanonicalized
values; update forward_with_edges (the sparse-edge entrypoint where
type_embedding is applied) to canonicalize charge_spin the same way forward()
does: if charge_spin is None substitute default_chg_spin, ensure
shape/broadcasting matches (nf, nloc, 2) or the expected per-atom shape, and
convert dtype/device to self.compute_dtype and extended_coord/device before
passing into _apply_charge_spin_embedding; alternatively, call the existing
normalization helper used by forward() from within forward_with_edges to avoid
duplicating logic.

---

Duplicate comments:
In `@deepmd/pt/model/descriptor/sezm_nn/so2.py`:
- Around line 865-868: The flag self.use_triton_rotations is set once in
__init__ using sezm_triton_enabled(device=self.device, dtype=self.dtype) and
then reused during forward rotations, causing device-mismatch after
model.to(...); change the branches in rotate_to_local_triton and
rotate_back_triton (and any other rotation call sites) to re-evaluate
sezm_triton_enabled at runtime using the actual tensor/device (e.g., use
x.device or the rotation tensor's device and dtype) instead of the cached
self.use_triton_rotations, and remove reliance on a frozen value (keep
sezm_triton_enabled import and call it where the Triton vs fallback branch is
selected so the correct backend is chosen after model.to(...)); also ensure
there is no _apply override that leaves the old flag stale.

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@deepmd/pt/entrypoints/freeze_pt2.py`:
- Around line 429-430: Update the docstring for the head parameter to reflect
current behavior: instead of stating "must be ``None``", document that the code
will auto-select "Default" for multitask checkpoints and otherwise requires an
explicit head name (or allow passing None to trigger auto-selection), and
mention the accepted values (e.g., "Default" or explicit head string) so callers
know the runtime semantics; target the docstring entry labeled head in
freeze_pt2.py (the docblock that currently says "Reserved for future multi-task
support; must be ``None``") and replace that sentence with the new, accurate
description.
- Around line 228-247: The metadata writer _collect_metadata() currently writes
an empty "type_map" (from model.get_type_map()) for SeZM/DPA4 models which leads
DeepEval._init_from_metadata() and DeepPotPTExpt to infer ntypes=0; fix
_collect_metadata() in freeze_pt2.py so that when list(model.get_type_map()) is
empty you populate a usable type map or explicit ntypes: e.g., set
metadata["type_map"] to a fallback list(range(N)) derived from
model.get_sel_type() or model.get_sel() (use len(model.get_sel_type()) or
len(model.get_sel()) to compute N), or include metadata["ntypes"]=int(N); ensure
the fallback logic uses model.get_sel_type()/model.get_sel() (and handles spin
if is_spin) so metadata-only loads give the correct ntypes for
DeepEval._init_from_metadata() and DeepPotPTExpt to consume.